import cn2an
import pandas as pd
import numpy as np
from pathlib import Path

import const

def read_and_clean(input_path: str) -> pd.DataFrame:
    """读取并清洗单个Excel文件"""
    # 读取第四个工作表（零部件汇总表）
    df = pd.read_excel(
        input_path,
        sheet_name=3,
        header=4,  # 第五行是标题
        skipfooter=1  # 跳过最后的总重统计行
    )
    return clean_dataframe(df)

def clean_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """数据清洗逻辑"""
    # 重命名列
    df.columns = [
        '序号', '图号或代号', '名称及规格', '单位', '数量', 
        '材料', '单重', '总重', '备注'
    ]
    
    # 拆分名称及规格, expand 表示将拆分结果扩展为两个独立的列
    df[['名称', '型号及规范']] = df['名称及规格'].str.split(
        r'\s+', n=1, expand=True
    ).fillna('/')

    # 处理显示格式
    df['材料'] = df['材料'].apply(lambda x: x if pd.notnull(x) else '/')
    df['数量'] = df['数量'].apply(lambda x: x if pd.notnull(x) else '/')
    df['单位'] = df['单位'].apply(lambda x: x if pd.notnull(x) else '/')
    df['单重'] = df['单重'].apply(lambda x: x if pd.notnull(x) else '/')
    df['总重'] = df['总重'].apply(lambda x: x if pd.notnull(x) else '/')
    df['型号及规范'] = df['型号及规范'].apply(lambda x: x if pd.notnull(x) else '/')
    
    # 移除临时列
    df.drop(columns=['序号', '图号或代号', '名称及规格'], inplace=True)
    
    return df

def merge_materials(df: pd.DataFrame) -> pd.DataFrame:
    """合并相同物料（处理整个目录数据）"""
    # 强制转换类型
    df['单重'] = pd.to_numeric(df['单重'], errors='coerce')
    df['数量'] = pd.to_numeric(df['数量'], errors='coerce')
    df['总重'] = pd.to_numeric(df['总重'], errors='coerce')
    
    group_cols = ['名称', '型号及规范', '材料', '单重']
    agg_rules = {
        '数量': 'sum',
        '总重': 'sum',
        '备注': lambda x: ';'.join(x.dropna().unique()),
        '单位': 'first',
        '单重': 'first'
    }
    
    merged = df.groupby(group_cols, as_index=False).agg(agg_rules)

    # 重新计算总重（双校验机制）
    merged['总重'] = merged.apply(
        lambda row: row['单重'] * row['数量'] 
        if (pd.notnull(row['单重']) and pd.notnull(row['数量'])) 
        else row['总重'],  # 保留原始总重
        axis=1
    )
    
    return merged

def merge_number(df: pd.DataFrame) -> pd.DataFrame:
    """处理单重为空, 数量不为空, 材料名称相同的, 把他们合并"""
    # 强制转换类型
    df['数量'] = pd.to_numeric(df['数量'], errors='coerce')

    group_cols = ['名称', '型号及规范', '材料', '单位', '单重', '总重']
    agg_rules = {
        '数量': 'sum',
        '备注': lambda x: ';'.join(x.dropna().unique()),
    }

    merged = df.groupby(group_cols, as_index=False).agg(agg_rules)

    return merged



def sort_and_group_materials(df: pd.DataFrame, sort_dict: dict) -> pd.DataFrame:
    """按分类排序并分组材料，支持尾部匹配，并添加从 1 开始的序号列"""
    # 创建一个新的列 "分类"，初始化为 NaN
    df['分类'] = np.nan
    # 在头部添加一个序号列, 初始化为 NaN
    df.insert(0, '序号', np.nan)
    
    # 根据 sort_dict 映射分类
    for category, materials in sort_dict.items():
        mask = df['名称'].apply(lambda x: any(x.endswith(keyword) for keyword in materials))
        df.loc[mask, '分类'] = category
    
    # 按分类排序（NaN 排到最后）
    df = df.sort_values(by='分类', na_position='last').reset_index(drop=True)
    
    # 生成序号列（从 1 开始）
    df['序号'] = df.groupby('分类', dropna=False).cumcount() + 1
    
    # 准备最终结果的容器
    result = []
    
    # 处理已分类的组和未分类的组
    for i, (category, group) in enumerate(df.groupby('分类', dropna=False), 1):
        # 创建标题行（仅名称列有值，其他列为 NaN）
        if pd.isna(category):
            header_row = pd.DataFrame([{'序号': cn2an.an2cn(i), '名称': '其他'}])
        else:
            header_row = pd.DataFrame([{'序号': cn2an.an2cn(i), '名称': category}])
        
        # 将标题行和当前组的数据合并到结果中
        result.append(header_row)
        result.append(group)
    
    # 合并所有结果
    final_df = pd.concat(result, ignore_index=True)
    
    # 删除临时的分类列（可根据需求保留）
    final_df = final_df.drop(columns=['分类'])
    
    return final_df

def format_columns(df: pd.DataFrame) -> pd.DataFrame:
    """格式化输出列"""
    column_order = [
        '名称', '型号及规范', '材料', '单位', 
        '数量', '单重', '总重', '备注'
    ]

    # 重新排列列顺序
    df = df.reindex(columns=column_order)

    # 分类
    df = sort_and_group_materials(df, const.sort_dict)

    return df
def save_to_excel(df: pd.DataFrame, output_path: str) -> None:
    """保存Excel文件"""
    with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
        df.to_excel(writer, index=False)

def process_all_data(combined_df: pd.DataFrame) -> pd.DataFrame:
    """处理合并后的完整数据集"""
    # 清理材料 名称 为空的行
    combined_df = combined_df[combined_df['名称'] != '/']

    # 分离 单重 为空的材料，单重为空，说明不用去计算总重
    combo_mask = (
        combined_df['单重'].isna() | 
        (combined_df['单重'] == '/') | 
        (combined_df['单重'] == '-') | 
        (combined_df['单重'] == '')
    )
    combo_df = combined_df[combo_mask].copy()
    normal_df = combined_df[~combo_mask].copy()
    
    # 合并普通材料
    merged_df = merge_materials(normal_df)
    # 合并单重为空的材料
    number_df = merge_number(combo_df)
    
    # 合并最终结果, ignore_index=True 表示忽略原始索引
    final_df = pd.concat([merged_df, number_df], ignore_index=True)
    return format_columns(final_df)

def process_directory(input_dir: str, output_dir: str) -> None:
    """处理目录中的所有Excel文件"""
    input_path = Path(input_dir)
    output_path = Path(output_dir)

    # 确保输出目录存在
    output_path.mkdir(parents=True, exist_ok=True)

    # 查找所有 .xlsx 文件
    xlsx_files = list(input_path.glob("*.xlsx"))
    xls_files = list(input_path.glob("*.xls"))

    # 合并所有文件
    excel_files = xlsx_files + xls_files

    all_dfs = []
    
    # 收集所有数据
    for file_path in excel_files:
        if not file_path.name.startswith('~$'):
            print(f"正在处理: {file_path.name}")
            try:
                df = read_and_clean(file_path)
                all_dfs.append(df)
            except Exception as e:
                print(f"  错误：{str(e)}")
    
    if not all_dfs:
        print("未找到有效Excel文件")
        return
    
    # 合并所有数据
    combined_df = pd.concat(all_dfs, ignore_index=True)
    
    # 统一处理合并后的数据
    final_df = process_all_data(combined_df)

    # 构建输出路径
    output_file = output_path / "combined.xlsx"
    
    # 保存结果
    save_to_excel(final_df, output_file)
    print(f"\n处理完成! 合并结果已保存至: {output_file}")